from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import Compose, ToTensor, Resize, Normalize
import os

model_path = './checkpoints/model.pth'
testpath = './datas/'
result_path = '/result/res.csv'

# source = [str(i) for i in range(0, 10)]
# source += [chr(i) for i in range(97, 97 + 26)]
source = [chr(i) for i in range(65, 65 + 26)]
alphabet = ''.join(source)


def img_loader(img_path):
    img = Image.open(img_path)
    return img.convert('RGB')


def make_dataset(data_path, alphabet, num_class, num_char):
    img_names = os.listdir(data_path)
    samples = []
    for img_name in img_names:
        img_path = os.path.join(data_path, img_name)
        target_str = img_name.split('.')[0]
        # assert len(target_str) == num_char
        target = []
        for char in target_str:
            vec = [0] * num_class
            vec[alphabet.find(char)] = 1
            target += vec
        samples.append((img_path, target))
    return samples


class CaptchaData(Dataset):
    def __init__(self,
                 data_path,
                 num_class=26,
                 num_char=4,
                 transform=None,
                 target_transform=None,
                 alphabet=alphabet):
        super(Dataset, self).__init__()
        self.data_path = data_path
        self.num_class = num_class
        self.num_char = num_char
        self.transform = transform
        self.target_transform = target_transform
        self.alphabet = alphabet
        self.samples = make_dataset(self.data_path, self.alphabet,
                                    self.num_class, self.num_char)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        img_path, target = self.samples[index]
        img = img_loader(img_path)
        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            target = self.target_transform(target)
        return img, torch.Tensor(target)


if __name__ == "__main__":
    batch_size = 16
    transforms = Compose([Resize((40, 120)), ToTensor()])

    train_dataset = CaptchaData('./datas/train/', transform=transforms)
    train_data_loader = DataLoader(train_dataset,
                                   batch_size=batch_size,
                                   num_workers=0,
                                   shuffle=True,
                                   drop_last=True)
    test_data = CaptchaData('./datas/test/', transform=transforms)
    test_data_loader = DataLoader(test_data,
                                  batch_size=batch_size,
                                  num_workers=0,
                                  shuffle=True,
                                  drop_last=True)
    for img, target in train_data_loader:
        print(img, target)